package com.jboltai.capability.vdb.types.milvus;

import com.alibaba.fastjson.JSONObject;
import com.jboltai.resource.embedding.EmbeddingModel;
import com.jboltai.resource.vdb.VDBResource;
import com.jboltai.resource.vdb.VDBType;
import io.milvus.client.MilvusServiceClient;
import io.milvus.param.R;
import io.milvus.param.RpcStatus;
import io.milvus.param.collection.CreateDatabaseParam;
import io.milvus.v2.client.ConnectConfig;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.service.database.request.CreateDatabaseReq;

public class MilvusVDBResource extends VDBResource {


    private MilvusClientV2 client;

    /**
     * 连接超时时间
     */
    private Long connectTimeout = null;

    /**
     * 用于设置保持连接的时间（以秒为单位）。如果在这段时间内没有任何请求发送到服务器，客户端将发送一个心跳包以保持连接。
     */
    private Long keepAliveTime = null;

    /**
     * 如果设置为 true，则即使没有任何调用，也会保持连接。如果设置为 false，则只有在有调用时才保持连接。
     */
    private Boolean aliveWithoutCalls = true;

    /**
     * 如果设置为 true，则使用安全的连接（例如 SSL/TLS）。如果设置为 false，则使用非安全的连接。
     */
    private Boolean secure = false;

    /**
     * 用于设置空闲超时时间（以秒为单位）。如果在这段时间内没有任何请求发送到服务器，客户端将关闭连接
     */
    private Long idleTimeout = null;




    /**
     * 资源
     * @param id 资源id
     * @param url
     * @param username
     * @param password
     */
    public MilvusVDBResource(String id, String url, String username, String password, EmbeddingModel embeddingModel) {
        super(id, VDBType.MILVUS, url, username, password,embeddingModel);
    }
    /**
     * 资源
     * @param id 资源id
     * @param url
     * @param username
     * @param password
     */
    public MilvusVDBResource(String id, String url, String username, String password, EmbeddingModel embeddingModel, JSONObject config) {
        super(id, VDBType.MILVUS, url, username, password,embeddingModel);
        init(config);
    }

    private void init( JSONObject config) {

        this.connectTimeout = getConfigValue(config, "connectTimeout", Long.class, 5l);
        this.keepAliveTime = getConfigValue(config, "keepAliveTime", Long.class,null);
        this.aliveWithoutCalls = getConfigValue(config, "aliveWithoutCalls", Boolean.class,true);
        this.secure = getConfigValue(config, "secure", Boolean.class, false);
        this.idleTimeout = getConfigValue(config, "idleTimeout", Long.class, null);


    }



    public MilvusClientV2 getClient() {
        if (client != null) {
            return client;
        }
        synchronized (this) {
            if (client != null) {
                return client;
            }
            ConnectConfig.ConnectConfigBuilder<?, ?> builder = ConnectConfig.builder();
            try {
                builder.uri(url).username(username)
                        .password(password);
                if (connectTimeout != null) {
                    builder.connectTimeoutMs(connectTimeout * 1000);
                }

                if (keepAliveTime != null) {
                    builder.keepAliveTimeMs(keepAliveTime * 1000);
                }
                if (aliveWithoutCalls != null) {
                    builder.keepAliveWithoutCalls(aliveWithoutCalls);
                }

                if (secure != null) {
                    builder.secure(secure);
                }
                if (idleTimeout != null) {
                    builder.idleTimeoutMs(idleTimeout * 1000);
                }
                builder.dbName(database);
                client = new MilvusClientV2(builder.build());

            } catch (Exception e) {
                if (e.getMessage().contains("database not found")) {
                    //新版Milvus，如果数据库不存在，调用    builder.withDatabaseName(database); 会报错
                    // 所以需要先通过Default将数据库创建出来
                    MilvusClientV2 defaultClient= null;
                    try {

                        CreateDatabaseReq req = CreateDatabaseReq.builder()
                                .databaseName(database)
                                .build();
                        defaultClient = getDefaultClient();
                        defaultClient.createDatabase(req);
                        client = new MilvusClientV2(builder.build());
                        return client;
                    } catch (Exception e2) {
                        throw new RuntimeException("Milvus向量数据库client链接失败，因为数据库["+database+"]不存在，且自动创建失败", e);
                    } finally {
                        if (defaultClient != null) {
                            defaultClient.close();
                        }
                    }
                } else {

                    throw new RuntimeException("Milvus向量数据库client链接失败", e);
                }
            }
        }
        return client;
    }

    /**
     * 创建数据库和检测数据库是否存在用的，因为 database还未创建好的情况下，调用builder.withDatabaseName(database);会报错
     * 所以通过使用default数据库，来完成创建
     * @return
     */
    private MilvusClientV2 getDefaultClient() {

        ConnectConfig.ConnectConfigBuilder<?, ?> builder = ConnectConfig.builder();
        try {
            builder.uri(url).username(username)
                    .password(password);
            if (connectTimeout != null) {
                builder.connectTimeoutMs(connectTimeout * 1000);
            }

            if (keepAliveTime != null) {
                builder.keepAliveTimeMs(keepAliveTime * 1000);
            }
            if (aliveWithoutCalls != null) {
                builder.keepAliveWithoutCalls(aliveWithoutCalls);
            }

            if (secure != null) {
                builder.secure(secure);
            }
            if (idleTimeout != null) {
                builder.idleTimeoutMs(idleTimeout * 1000);
            }
            builder.dbName("default");
            client = new MilvusClientV2(builder.build());

        } catch (Exception e) {
            throw new RuntimeException("Milvus向量数据库client链接失败", e);
        }
        return client;
    }

    public Long getConnectTimeout() {
        return connectTimeout;
    }

    /**
     * 设置连接超时时间
     * @param connectTimeout
     * @return
     */
    public MilvusVDBResource setConnectTimeout(Long connectTimeout) {
        this.connectTimeout = connectTimeout;
        return this;
    }

    public Long getKeepAliveTime() {
        return keepAliveTime;
    }

    /**
     * 设置保持连接的时间（以秒为单位）。如果在这段时间内没有任何请求发送到服务器，客户端将发送一个心跳包以保持连接。
     * @param keepAliveTime
     * @return
     */
    public MilvusVDBResource setKeepAliveTime(Long keepAliveTime) {
        this.keepAliveTime = keepAliveTime;
        return this;
    }

    public Boolean getAliveWithoutCalls() {
        return aliveWithoutCalls;
    }

    /**
     * 设置是否即使没有任何调用，也会保持连接。如果设置为 false，则只有在有调用时才保持连接。
     * @param aliveWithoutCalls
     * @return
     */
    public MilvusVDBResource setAliveWithoutCalls(Boolean aliveWithoutCalls) {
        this.aliveWithoutCalls = aliveWithoutCalls;
        return this;
    }

    public Boolean getSecure() {
        return secure;
    }

    /**
     * 设置是否使用安全的连接（例如 SSL/TLS）。如果设置为 false，则使用非安全的连接。
     * @param secure
     * @return
     */
    public MilvusVDBResource setSecure(Boolean secure) {
        this.secure = secure;
        return this;
    }

    public Long getIdleTimeout() {
        return idleTimeout;
    }

    /**
     * 设置空闲超时时间（以秒为单位）。如果在这段时间内没有任何请求发送到服务器，客户端将关闭连接
     * @param idleTimeout
     * @return
     */
    public MilvusVDBResource setIdleTimeout(Long idleTimeout) {
        this.idleTimeout = idleTimeout;
        return this;
    }
}
